"""
    Title

    author: wxz
    date: 2021-12-13
    github: https://github.com/xinzwang
"""

import numpy as np


class CMA_ES(object):
    """
    CMA-ES 协方差自适应进化
    """

    def __init__(self, loss_func, n):
        """设置算法超参数

        协方差自适应参数采用论文中默认数值。这些参数一般具有较好鲁棒性，不需要针对问题进行修改
        lam(子代数量)、mu(父代数量)、w(重组权重)相对而言不是很重要，可以根据需要进行调整
        """
        # 待优化函数指针
        self.loss_func = loss_func
        # 迭代次数
        self.EPOCH = 100
        # 解空间维度
        self.n = n
        # 子代数量    每一代采样个体数
        self.lam = 4 + np.int32(np.floor(3 * np.log(n)))  # lam = 4 + floor(3 * ln(n))
        # 父代数量    取前mu的样本
        mu_dot = self.lam / 2
        self.mu = np.int32(np.floor(mu_dot))
        # 重组权值
        w_dot = np.log(mu_dot + 0.5) - np.log(np.linspace(1, self.mu, self.mu))
        self.w = w_dot / np.sum(w_dot)
        self.mu_eff = 1.0 / (np.sum(np.square(self.w)))  # 方差有效选择质量
        # 步长控制参数
        self.c_sigma = (self.mu_eff + 2) / (n + self.mu_eff + 5)
        self.d_sigma = 1 + 2 * np.max([0, np.sqrt((self.mu_eff - 1) / (n + 1)) - 1]) + self.c_sigma
        # 协方差自适应参数
        alpha_mu = 2
        self.c_1 = 2.0 / (np.square(n + 1.3) + self.mu_eff)
        self.c_c = (4 + self.mu_eff / n) / (n + 4 + 2 * self.mu_eff / n)
        self.c_mu = min(1 - self.c_1, alpha_mu * (self.mu_eff - 2 + 1.0 / self.mu_eff) / (
                np.square(n + 2) + alpha_mu * self.mu_eff / 2))
        # 其它参数
        self.E_N = np.sqrt(self.n)  # E||N(0,I)|| 服从标准正太分布的变量的二范式的期望

        """初始化参数"""
        # 样本
        self.x = None
        self.y = None
        self.z = None
        self.y_w = None  # 加权后的y
        # 样本估值
        self.loss = None
        # 协方差矩阵     # TODO: 协方差矩阵初始化，可以改进
        self.C = np.identity(self.n, dtype=np.float32)  # 初始化为单位矩阵I
        # self.C = np.ones([self.n, self.n], dtype=np.float32)
        self.B = None
        self.D = None
        # 进化路径
        self.p_sigma = 0
        self.p_c = 0
        # 进化代数
        self.g = 0
        # 均值
        self.m = 0  # 根据问题选择
        # 步长
        self.sigma = 1  # 根据问题选择

        return

    def heaviside(self):
        """
        赫维赛德函数
        """
        g = self.g
        p_sigma = self.p_sigma
        c_sigma = self.c_sigma
        val1 = np.sum(np.square(p_sigma)) / np.sqrt(1 - np.power(1 - c_sigma, 2 * (g + 1)))
        val2 = (1.4 + 2 / (self.n + 1)) * self.E_N
        out = 1 if val1 < val2 else 0
        return out

    def cov_mu(self):
        c_mu = self.c_mu
        w = self.w
        y = self.y
        # 计算
        index = np.argsort(self.loss)
        index = index[0:self.mu]
        y_mu = self.y[index]  # 前mu优秀的y集合
        out = c_mu * np.sum(w.dot(y_mu).dot(y_mu.T))
        return out

    def sample_population(self):
        """
        根据参数，采样计算新的参数样本
        """
        # 由于协方差矩阵C是对称矩阵
        # 因此D(特征向量矩阵)是正交矩阵
        # 即: D.I = D.T
        D, B = np.linalg.eig(self.C)  # 特征分解  D是特征值  B是特征向量的标准正交基
        D = np.diag(np.sqrt(D))  # D是正特征值的平方根

        # 采样: z ~ N(0, I)    lam个
        z = np.random.multivariate_normal(mean=np.zeros(self.n), cov=np.eye(self.n), size=self.lam)

        # 计算: y = B*D*z ~ N(0, C)
        y = np.zeros_like(z)
        A = D.dot(B)
        for i in range(self.lam):
            y[i, :] = (A.dot(z[i, :].T)).T

        # 计算: x = m + sigma*y ~ N(m,sigma^2*C)
        x = self.m + self.sigma * y

        # 计算: 新参数的loss
        loss = self.loss_func(x)

        return (x, y, z, loss), (D, B)

    def update_m(self):
        """
        计算均值
        """
        # 从种群中获取loss前 mu的个体作为新的父代
        # index  是父代样本的索引 按loss从小到大排序
        # y_mu   是父代样本集合
        index = np.argsort(self.loss, axis=0)
        index = index[0:self.mu]
        y_mu = self.y[index]  # 前mu优秀的y集合

        # 按照权值对父代样本进行加权
        # y_w  是加权后的样本向量，即期望
        y_w = (self.w).dot(y_mu)  # y_w = sum(w*y)

        # 按照公式更新均值
        # 这里m应是x的期望。而x是由y进行线性变换而产生
        # 因此将y的期望做线性变换，从而正确得到x的期望
        m = self.m + self.sigma * y_w  # m = m + sigma*y_w

        # 返回：新的均值、y的期望
        return m, y_w

    def update_sigma(self):
        """
        计算更新步长
        """
        # 解析存储的各变量
        B = self.B
        D = self.D
        sigma = self.sigma
        c_sigma = self.c_sigma
        d_sigma = self.d_sigma
        mu_eff = self.mu_eff
        y_w = self.y_w
        E_N = self.E_N

        # 更新sigma路径
        p_sigma = (1 - c_sigma) * self.p_sigma + np.sqrt(c_sigma * (2 - c_sigma) * mu_eff) * (
            B.dot(np.linalg.inv(D)).dot(B.T)) * y_w

        sigma = sigma * np.exp(c_sigma / d_sigma * np.sum(np.square(p_sigma)) / E_N - 1)

        return sigma, p_sigma

    def update_C(self):
        c_1 = self.c_1
        c_c = self.c_c
        c_mu = self.c_mu
        p_c = self.p_c
        mu_eff = self.mu_eff
        y_w = self.y_w
        C = self.C

        # 计算h_sigma
        h_sigma = self.heaviside()
        delta_h_sigma = (1 - h_sigma) * c_c * (2 - c_c)

        p_c = (1 - c_c) * p_c + h_sigma * np.sqrt(c_c * (2 - c_c) * mu_eff) * y_w
        C = (1 - c_1 - c_mu) * C + c_1 * (p_c * p_c.T + delta_h_sigma * C) + self.cov_mu()

        return p_c, C

    def solve(self):
        while self.g < self.EPOCH:  # 满足终止条件
            self.g += 1

            # 1.采样新的人口
            (self.x, self.y, self.z, self.loss), (self.D, self.B) = self.sample_population()

            # 2.计算均值 m
            self.m, self.y_w = self.update_m()

            # 3.计算步长 sigma
            self.sigma, self.p_sigma = self.update_sigma()

            # 4.更新协方差矩阵 C
            self.p_c, self.C = self.update_C()

            print("epoch=%d    loss=%.5f" % (self.g, np.average(self.loss)),
                  "    m={}  sigma={}".format(self.m, self.sigma))
        return self.x


if __name__ == '__main__':
    from funcs import *

    DIM = 2

    f = F_Rosenbrock(DIM)

    alg = CMA_ES(f.forward, n=DIM)

    x = alg.solve()

    pass
